import numpy as np
from scipy.sparse import csr_matrix, vstack
from common.img_utils.sas_python.utils import (
    makeweights,
    prune_knn,
    adjacency,
    sparse,
)
from common.img_utils.sas_python.utils import vector_idx_minus_1_from_matlab


def build_bipartite_graph_for_matlab(img_size, para, seg, seg_vals, seg_edges):
    Np = img_size[0] * img_size[1]

    # Get the overall number of superpixels
    Nsp = 0
    for k in range(0, len(seg)):
        Nsp = Nsp + len(seg[k][0])

    W_Y = csr_matrix((Nsp, Nsp))
    edgesXY = None
    j = 0
    for k in range(len(seg)):  # For each over-segmentation
        # Affinities between superpixels
        w = makeweights(seg_edges[k], seg_vals[k], para.beta)
        W = adjacency(seg_edges[k], w)
        Nk = len(seg[k][0])  # Number of superpixels in over - segmentation k
        W_Y[j : j + Nk, j : j + Nk] = prune_knn(W, para.nb)

        # Affinities between pixels and superpixels
        for i in range(Nk):
            idxp = seg[k][0][i]  # pixel indices in superpixel i
            idxp_idxes = vector_idx_minus_1_from_matlab(idxp)
            Nki = len(idxp_idxes)
            idxsp = j + np.zeros(Nki)
            idxsp = np.expand_dims(idxsp, axis=1)
            idx_p_sp = np.concatenate((idxp_idxes, idxsp), axis=1)
            if edgesXY is None:
                edgesXY = idx_p_sp.copy()
            else:
                edgesXY = np.append(edgesXY, idx_p_sp, axis=0)
            j = j + 1

    data = np.asarray([para.alpha] * edgesXY.shape[0])
    W_XY = sparse(edgesXY[:, 0], edgesXY[:, 1], data, Np, Nsp)

    # Affinity between a superpixel and itself is set to be the maximum 1.
    W_Y.setdiag(1)
    B = vstack((W_XY, W_Y))

    return B


def build_bipartite_graph(img_size, para, seg, seg_vals, seg_edges):
    """Build the bipartite graph

    Reference:
    Li, Z., Wu, X. M., & Chang, S. F. (2012, June).
    Segmentation using superpixels: A bipartite graph partitioning approach.
    In 2012 IEEE conference on computer vision and pattern recognition (pp. 789-796). IEEE.
    https://github.com/ColumbiaDVMM/Segmentation-Using-Superpixels/blob/master/algorithms/build_bipartite_graph.m

    Parameters
    ----------
    img_size : list
        [h, w]
    para : a SASPara entity
        need para.alpha, para.beta, para.nb
        alpha : affinity between pixels and superpixels [float]
        beta : scale factor in superpixel affinity [int]
        nb: number of neighbors for superpixels [int]
    seg : list
        list of dictionaries. Each dict has the length of (# of superpixel).
        In each key, the values are pixels' indices which belong to this superpixel.
    seg_vals : list
        list of list. Each has the length of (# of superpixel).
        In each element, the values are this superpixel's mean BGR(N) values or RGB(N) values.
    seg_edges : list
        list of list. A ndarray include all superpixels' edges ((# of edge pairs) x 2)

    Returns
    -------
    B : ndarray
        a bipartite graph
    """
    Np = img_size[0] * img_size[1]

    # Get the overall number of superpixels
    Nsp = 0
    for k in range(len(seg)):
        Nsp = Nsp + len(seg[str(k)])

    W_Y = csr_matrix((Nsp, Nsp))
    edgesXY = None
    j = 0
    for k in range(len(seg)):  # For each over-segmentation
        print("Process sp labels {}...".format(k + 1))

        if (
            len(seg[str(k)]) == 0
            and len(seg_vals[str(k)]) == 0
            and len(seg_edges[str(k)]) == 0
        ):
            continue

        # Affinities between superpixels
        w = makeweights(seg_edges[str(k)], seg_vals[str(k)], para.beta)
        W = adjacency(seg_edges[str(k)], w, N=len(seg[str(k)]))
        Nk = len(seg[str(k)])  # Number of superpixels in over - segmentation k
        W_Y[j : j + Nk, j : j + Nk] = prune_knn(W, para.nb)

        # Affinities between pixels and superpixels
        for i in range(Nk):
            idxp = seg[str(k)][str(i)]  # pixel indices in superpixel i
            Nki = len(idxp)
            idxsp = j + np.zeros(Nki)
            idxp = np.expand_dims(idxp, axis=1)
            idxsp = np.expand_dims(idxsp, axis=1)
            idx_p_sp = np.concatenate((idxp, idxsp), axis=1)
            if edgesXY is None:
                edgesXY = idx_p_sp.copy()
            else:
                edgesXY = np.append(edgesXY, idx_p_sp, axis=0)
            j = j + 1

    data = np.asarray([para.alpha] * edgesXY.shape[0])
    W_XY = sparse(edgesXY[:, 0], edgesXY[:, 1], data, Np, Nsp)

    # Affinity between a superpixel and itself is set to be the maximum 1.
    W_Y.setdiag(1)
    B = vstack((W_XY, W_Y))

    return B
